{
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  },
  "name": "",
  "signature": "sha256:d80f947f486394d9a39653ee4fed22be322470c3fa0623d83638154389115728"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Detection with SSD\n",
      "\n",
      "In this example, we will load a SSD model and use it to detect humans."
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "### 1. Setup\n",
      "\n",
      "* First, Load necessary libs and set up caffe and caffe_root.\n",
      "Please run this file in caffe_root"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy as np\n",
      "import matplotlib.pyplot as plt\n",
      "%matplotlib inline\n",
      "\n",
      "plt.rcParams['figure.figsize'] = (10, 10)\n",
      "plt.rcParams['image.interpolation'] = 'nearest'\n",
      "plt.rcParams['image.cmap'] = 'gray'\n",
      "\n",
      "# Make sure that caffe is on the python path:\n",
      "caffe_root = '../../'  # this file is expected to be in {caffe_root}/examples/mppp\n",
      "import os\n",
      "os.chdir(caffe_root)\n",
      "import sys\n",
      "sys.path.insert(0, 'python')\n",
      "\n",
      "import caffe\n",
      "caffe.set_device(0)\n",
      "caffe.set_mode_gpu()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "* Load the net in the test phase for inference, and configure input preprocessing.\n",
      "\u628a\u8fd9\u91cc\u7684model_def \u548c model_weights\u6539\u6210\u6211\u4eec\u8981\u7528\u7684\u90a3\u4e2a"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "model_def = 'models/VGG_SSD/deploy.prototxt'\n",
      "model_weights = 'models/VGG_SSD/VGG_MPII_COCO14_SSD_500x500_iter_60000.caffemodel'\n",
      "\n",
      "net = caffe.Net(model_def,      # defines the structure of the model\n",
      "                model_weights,  # contains the trained weights\n",
      "                caffe.TEST)     # use test mode (e.g., don't perform dropout)\n",
      "\n",
      "# input preprocessing: 'data' is the name of the input blob == net.inputs[0]\n",
      "transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})\n",
      "transformer.set_transpose('data', (2, 0, 1))\n",
      "transformer.set_mean('data', np.array([104,117,123])) # mean pixel\n",
      "transformer.set_raw_scale('data', 255)  # the reference model operates on images in [0,255] range instead of [0,1]\n",
      "transformer.set_channel_swap('data', (2,1,0))  # the reference model has channels in BGR order instead of RGB"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "heading",
     "level": 3,
     "metadata": {},
     "source": [
      "2. SSD detection"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "* Load all image.\n",
      "image_resize to 500,\n",
      "root_dir is where MPII images stored,\n",
      "configThred the threshold for bounding box"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# set net to batch size of 1\n",
      "import os\n",
      "import h5py\n",
      "image_resize = 500\n",
      "net.blobs['data'].reshape(1,3,image_resize,image_resize)\n",
      "root_dir = \"data/MPII/images\"\n",
      "\n",
      "configThred = 0.09\n",
      "NMSThred = 0.45\n",
      "\n",
      "lines = [line.rstrip('\\n') .rstrip('\\r') for line in open(\"examples/rmpe/util/group_test_list.txt\")]\n",
      "\n",
      "directory = 'examples/rmpe/mpii-test'+str(configThred)+'/'\n",
      "if not os.path.exists(directory):\n",
      "    os.makedirs(directory)\n",
      "results = open(directory+\"test-bbox_images.txt\", 'w')\n",
      "score_file = open(directory+\"score.txt\",'w')\n",
      "index_file = open(directory+\"index.txt\",'w')\n",
      "\n",
      "FileLength = len(lines)\n",
      "\n",
      "num_boxes=0\n",
      "\n",
      "xminarr=[]\n",
      "yminarr=[]\n",
      "xmaxarr=[]\n",
      "ymaxarr=[]\n",
      "\n",
      "for i in xrange(FileLength):\n",
      "    if ((i%1000) == 0):\n",
      "        print i\n",
      "    picture = lines[i].split(\"\\t\")\n",
      "    p_name = picture[0]\n",
      "\n",
      "    filename = os.path.join(root_dir, p_name)\n",
      "\n",
      "    if (os.path.isfile(filename) == False):\n",
      "        print filename + \" does not exist\"\n",
      "        continue\n",
      "    image = caffe.io.load_image(filename)\n",
      "\n",
      "    transformed_image = transformer.preprocess('data', image)\n",
      "    net.blobs['data'].data[...] = transformed_image\n",
      "    \n",
      "    # Forward pass.\n",
      "    detections = (net.forward()['detection_out'])\n",
      "\n",
      "    # Parse the outputs.\n",
      "    det_label = detections[0,0,:,1]\n",
      "    det_conf = detections[0,0,:,2]\n",
      "    det_xmin = detections[0,0,:,3]\n",
      "    det_ymin = detections[0,0,:,4]\n",
      "    det_xmax = detections[0,0,:,5]\n",
      "    det_ymax = detections[0,0,:,6]\n",
      "\n",
      "    top_indices = [m for m, conf in enumerate(det_conf) if conf > configThred]\n",
      "    top_conf = det_conf[top_indices]\n",
      "    top_label_indices = det_label[top_indices].tolist()\n",
      "    top_labels = det_label[top_indices]\n",
      "    top_xmin = det_xmin[top_indices]\n",
      "    top_ymin = det_ymin[top_indices]\n",
      "    top_xmax = det_xmax[top_indices]\n",
      "    top_ymax = det_ymax[top_indices]\n",
      "\n",
      "    colors = ['b', 'g', 'r', 'c', 'm', 'y', 'k']\n",
      "\n",
      "    if(top_conf.shape[0]!=0):\n",
      "        index_file.write(\"{} {} \".format(p_name,num_boxes+1))\n",
      "    for k in xrange(top_conf.shape[0]):\n",
      "        label = top_labels[k]\n",
      "        if (label != 1):\n",
      "            continue\n",
      "        xmin = int(round(top_xmin[k] * image.shape[1]))\n",
      "        ymin = int(round(top_ymin[k] * image.shape[0]))\n",
      "        xmax = int(round(top_xmax[k] * image.shape[1]))\n",
      "        ymax = int(round(top_ymax[k] * image.shape[0]))\n",
      "        score = top_conf[k]\n",
      "        if (xmin <= 0):\n",
      "            xmin = 1\n",
      "        if (ymin <= 0):\n",
      "            ymin = 1\n",
      "        xminarr.append(xmin);yminarr.append(ymin);xmaxarr.append(xmax);ymaxarr.append(ymax);\n",
      "        results.write(\"{}\\n\".format(p_name))\n",
      "        score_file.write(\"{}\\n\".format(score))\n",
      "\n",
      "        num_boxes += 1\n",
      "    \n",
      "    if(top_conf.shape[0]!=0):\n",
      "        index_file.write(\"{}\\n\".format(num_boxes))\n",
      "print \"Average Boxes per image:\", float(num_boxes)/FileLength\n",
      "results.close()    \n",
      "score_file.close()\n",
      "index_file.close()\n",
      "with h5py.File(directory+'test-bbox.h5', 'w') as hf:\n",
      "                hf.create_dataset('xmin', data=np.array(xminarr))\n",
      "                hf.create_dataset('ymin', data=np.array(yminarr))\n",
      "                hf.create_dataset('xmax', data=np.array(xmaxarr))\n",
      "                hf.create_dataset('ymax', data=np.array(ymaxarr))\n",
      "print \"Done\""
     ],
     "language": "python",
     "metadata": {
      "scrolled": true
     },
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "0\n",
        "1000"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "Average Boxes per image:"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        " 8.29635949943\n",
        "Done\n"
       ]
      }
     ],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print num_boxes"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "14585\n"
       ]
      }
     ],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": null
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": null
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": null
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": null
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": null
    }
   ],
   "metadata": {}
  }
 ]
}